from ultralytics import YOLO


def export_to_onnx(model_path, output_path):
    """导出 YOLOv8 模型为优化 ONNX 格式"""
    model = YOLO(model_path)

    # 关键优化参数
    model.export(
        format='onnx',
        imgsz=640,  # 固定输入尺寸
        dynamic=False,  # 静态批次更高效
        simplify=True,  # 启用模型简化
        opset=17,  # 使用最新算子集
        half=True,  # FP16 量化
        batch=1,  # 指定批次大小
        device='cuda',  # 使用 GPU 加速导出
        workspace=4,  # 4GB 显存
        nms=True,  # 包含 NMS 层
    )
    print(f"✅ ONNX 模型导出成功: {output_path}")


# 示例使用
export_to_onnx('../models/yolov8n.pt', '../models/yolov8n_optimized.onnx')